from flask import Blueprint, request, jsonify
from flask_jwt_extended import jwt_required, get_jwt_identity
from models import db, Problem, User
from services.judge_service import JudgeService
from services.async_judge_service import async_judge_service
from decorators import student_required
from utils import format_response
import json

judge_bp = Blueprint('judge', __name__)
judge_service = JudgeService()

@judge_bp.route('/run', methods=['POST'])
@jwt_required()
@student_required
def run_code():
    """运行代码（不保存提交记录）"""
    try:
        data = request.get_json()
        problem_id = data.get('problem_id')
        language = data.get('language')
        code = data.get('code')
        
        if not all([problem_id, language, code]):
            return jsonify(format_response(message="缺少必要参数", code=400)), 400
        
        # 获取题目信息
        problem = Problem.query.get(problem_id)
        if not problem:
            return jsonify(format_response(message="题目不存在", code=404)), 404
        
        # 解析测试用例
        test_cases = []
        if problem.test_cases:
            try:
                test_cases = json.loads(problem.test_cases)
            except json.JSONDecodeError:
                test_cases = []
        
        # 运行代码评测
        judge_result = judge_service.judge_submission(
            code=code,
            language=language,
            test_cases=test_cases,
            time_limit=problem.time_limit,
            memory_limit=problem.memory_limit
        )
        
        return jsonify(format_response(judge_result, "代码运行完成"))
        
    except Exception as e:
        return jsonify(format_response(message=str(e), code=500)), 500

@judge_bp.route('/submit', methods=['POST'])
@jwt_required()
@student_required
def submit_code():
    """提交代码（保存提交记录）"""
    try:
        data = request.get_json()
        problem_id = data.get('problem_id')
        language = data.get('language')
        code = data.get('code')
        
        if not all([problem_id, language, code]):
            return jsonify(format_response(message="缺少必要参数", code=400)), 400
        
        # 获取题目信息
        problem = Problem.query.get(problem_id)
        if not problem:
            return jsonify(format_response(message="题目不存在", code=404)), 404
        
        # 解析测试用例
        test_cases = []
        if problem.test_cases:
            try:
                test_cases = json.loads(problem.test_cases)
            except json.JSONDecodeError:
                test_cases = []
        
        # 运行代码评测
        judge_result = judge_service.judge_submission(
            code=code,
            language=language,
            test_cases=test_cases,
            time_limit=problem.time_limit,
            memory_limit=problem.memory_limit
        )
        
        # 计算得分
        score = judge_result.get('score', 0)
        if judge_result['status'] == 'AC':
            score = problem.points
        elif judge_result['status'] == 'WA':
            # 部分得分：根据通过的测试用例数量计算
            passed_cases = sum(1 for case in judge_result.get('test_cases', []) if case['status'] == 'AC')
            total_cases = len(judge_result.get('test_cases', []))
            if total_cases > 0:
                score = int((passed_cases / total_cases) * problem.points)
        
        # 保存提交记录（这里可以扩展为保存到数据库）
        # TODO: 实现提交记录保存
        
        return jsonify(format_response({
            'result': judge_result,
            'score': score,
            'max_score': problem.points
        }, "代码提交完成"))
        
    except Exception as e:
        return jsonify(format_response(message=str(e), code=500)), 500

# 新增异步判题相关接口
@judge_bp.route('/submit-async', methods=['POST'])
@jwt_required()
@student_required
def submit_code_async():
    """异步提交代码"""
    try:
        data = request.get_json()
        problem_id = data.get('problem_id')
        language = data.get('language')
        code = data.get('code')
        
        if not all([problem_id, language, code]):
            return jsonify(format_response(message="缺少必要参数", code=400)), 400
        
        # 获取当前用户ID
        user_id = get_jwt_identity()
        
        # 获取题目信息
        problem = Problem.query.get(problem_id)
        if not problem:
            return jsonify(format_response(message="题目不存在", code=404)), 404
        
        # 解析测试用例
        test_cases = []
        if problem.test_cases:
            try:
                test_cases = json.loads(problem.test_cases)
            except json.JSONDecodeError:
                test_cases = []
        
        # 提交异步判题任务
        task_id = async_judge_service.submit_judge_task(
            user_id=user_id,
            problem_id=problem_id,
            code=code,
            language=language,
            test_cases=test_cases,
            time_limit=problem.time_limit,
            memory_limit=problem.memory_limit
        )
        
        return jsonify(format_response({
            'task_id': task_id,
            'status': 'pending'
        }, "判题任务已提交"))
        
    except Exception as e:
        return jsonify(format_response(message=str(e), code=500)), 500

@judge_bp.route('/status/<task_id>', methods=['GET'])
@jwt_required()
@student_required
def get_judge_status(task_id):
    """获取判题状态"""
    try:
        status = async_judge_service.get_task_status(task_id)
        if not status:
            return jsonify(format_response(message="任务不存在", code=404)), 404
        
        return jsonify(format_response(status, "获取状态成功"))
        
    except Exception as e:
        return jsonify(format_response(message=str(e), code=500)), 500

@judge_bp.route('/queue-status', methods=['GET'])
@jwt_required()
def get_queue_status():
    """获取队列状态"""
    try:
        status = async_judge_service.get_queue_status()
        return jsonify(format_response(status, "获取队列状态成功"))
        
    except Exception as e:
        return jsonify(format_response(message=str(e), code=500)), 500
